# waitress-serve --listen=127.0.0.1:1234 Qwn:app
import cachetools
from flask import Flask, jsonify, request
from flask_cors import CORS
from waitress import serve
import requests
import json
global context
app = Flask(__name__)
CORS(app)

# 创建一个带有TTL过期时间的缓存，例如：最多存储XX个条目，每条记录有效期为XX分钟
context_cache = cachetools.TTLCache(maxsize=1, ttl=60*5)

@app.route('/', methods=['GET'])
def index():
    return jsonify({'message': 'Welcome to Qwn API'})


@app.route('/', methods=['POST'])
def post_data():
    data = request.json
    response_ollama_def = ollama_def(data.get('prompt', ''))
    return jsonify({'response': response_ollama_def})


def ollama_def(data="你是谁？，能干什么？"):
    url = "http://localhost:11434/api/generate"
    payload = {
        "model": "qwen2.5-coder:7b",
        "prompt": data,
        "stream": False,

    }
    # 检查是否有缓存的context，并添加到payload中
    context = context_cache.get('context')
    if context:
        payload['context'] = context
    response = requests.post(url, json=payload)
    print('发送qwn请求：   ' + f"{payload}")
    print('响应qwn请求：   ' + f"{response.text}")
    obj = json.loads(response.text.strip())

    new_context = obj.get('context')
    if new_context is not None:
        context_cache['context'] = new_context
    obj_response = obj.get('response')
    return obj_response


print('INFO  waitress服务器在  http://127.0.0.1:1234  启动')
if __name__ == '__main__':
    # ollama_def()
    serve(app, host='127.0.0.1', port=1234)
